# 增加 pygame 库引用

import pygame
from typing import Dict, List, Tuple
from algo_struct.algo import rad_limit
import os
# 仿真器入口 用于训练强化学习 跟踪算法用的简化版本仿真器
# 由于没有GUI界面 
# 航控需要通过 位置偏差 和 航向偏差 等指标查看控制算法是否有效 
from dynamics_simulation import motionSim
import numpy as np
import matplotlib.pyplot as plt
import socket
import json
from navigation import gpsAndUtm
from collections import namedtuple, OrderedDict
import traceback
from algo_struct.algo import rad_limit
import pickle, os
import math


# 解析控制参数
def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    delta = recv_obser['rudl']
    rsp = recv_obser['rspl']

    return delta, rsp

# 打包船的状态返回
def encode_data(trans):
    # 按照标准接口编写
    data = json.dumps(trans, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')

# 计算奖励
def get_reward(delta_psi, dpsi, delta):
    # DDPG高频大舵角 得分高 但是这是执行机构做不到的 因此增加累积舵角惩罚
    score = -((0.1*(delta_psi)**2) + (dpsi**2) + (delta**2))

    print(f'score is {score}')
    return score

def distance_point_to_line(x0, y0, x1, y1, x2, y2):
    numerator = np.abs((y2-y1)*x0 - (x2-x1)* y0 + x2*y1 - y2* x1)
    denominator = math.sqrt((y2 - y1)**2 + (x2 - x1)**2)
    return numerator / denominator

# Hybrid A*算法封装
def hybrid_a_star(start, goal, heading, grid_map, resolution=1.0, max_iterations=1000):
    """
    实现Hybrid A*算法
    :param start: 起点坐标 (x, y)
    :param goal: 终点坐标 (x, y)
    :param heading: 起点朝向 (角度)
    :param grid_map: 网格地图 (二维数组，0表示可通行，1表示障碍物)
    :param resolution: 网格分辨率
    :param max_iterations: 最大迭代次数
    :return: 规划的路径 (path_x, path_y)
    """
    from queue import PriorityQueue

    # 定义节点类
    class Node:
        def __init__(self, x, y, theta, cost, parent=None):
            self.x = x
            self.y = y
            self.theta = theta
            self.cost = cost
            self.parent = parent

        def __lt__(self, other):
            return self.cost < other.cost

    # 初始化起点和终点
    start_node = Node(start[0], start[1], heading, 0)
    goal_node = Node(goal[0], goal[1], 0, 0)

    # 定义动作集 (线速度, 角速度)
    actions = [(1, 0), (1, np.pi/6), (1, -np.pi/6)]

    # 初始化优先队列和已访问节点集合
    open_set = PriorityQueue()
    open_set.put(start_node)
    closed_set = set()

    # A*算法主循环
    iteration = 0
    while not open_set.empty() and iteration < max_iterations:
        current_node = open_set.get()

        # 检查是否到达目标点
        if np.sqrt((current_node.x - goal_node.x)**2 + (current_node.y - goal_node.y)**2) < resolution:
            path_x, path_y = [], []
            while current_node is not None:
                path_x.append(current_node.x)
                path_y.append(current_node.y)
                current_node = current_node.parent
            return path_x[::-1], path_y[::-1]  # 反转路径

        # 扩展节点
        for action in actions:
            new_x = current_node.x + action[0] * np.cos(current_node.theta)
            new_y = current_node.y + action[0] * np.sin(current_node.theta)
            new_theta = current_node.theta + action[1]

            # 检查新节点是否在地图范围内且无障碍物
            if 0 <= new_x < len(grid_map) and 0 <= new_y < len(grid_map[0]) and grid_map[int(new_x)][int(new_y)] == 0:
                new_cost = current_node.cost + action[0]  # 简化代价计算
                new_node = Node(new_x, new_y, new_theta, new_cost, current_node)

                # 检查是否已访问过
                if (new_x, new_y) not in closed_set:
                    open_set.put(new_node)
                    closed_set.add((new_x, new_y))

        iteration += 1

    # 如果未找到路径，返回空路径
    return [], []

# 平滑路径函数
def smooth_path(path_x, path_y, weight_data=0.5, weight_smooth=0.1, tolerance=0.000001):
    new_path = [[float(path_x[i]), float(path_y[i])] for i in range(len(path_x))]
    change = tolerance
    while change >= tolerance:
        change = 0.0
        for i in range(1, len(path_x)-1):
            for j in range(len(path_x[0])):
                aux = new_path[i][j]
                new_path[i][j] += weight_data * (path_x[i][j] - new_path[i][j])
                new_path[i][j] += weight_smooth * (new_path[i-1][j] + new_path[i+1][j] - (2.0 * new_path[i][j]))
                change += abs(aux - new_path[i][j])
    return [p[0] for p in new_path], [p[1] for p in new_path]

# 新增绘制路径的函数
def draw_path(screen, path_x, path_y, color=(0, 0, 255), width=2):
    for i in range(len(path_x) - 1):
        pygame.draw.line(screen, color, (path_x[i], path_y[i]), (path_x[i + 1], path_y[i + 1]), width)

# 新增绘制障碍物的函数
def draw_obstacles(screen, grid_map, resolution, color=(0, 0, 0)):
    for i in range(len(grid_map)):
        for j in range(len(grid_map[0])):
            if grid_map[i][j] == 1:
                rect = pygame.Rect(i * resolution, j * resolution, resolution, resolution)
                pygame.draw.rect(screen, color, rect)

# test motion sim
if __name__ == '__main__':
    # 与航行控制算法连接TCP
    LOCAL_HOST = "127.0.0.1"
    LOCAL_PORT = 10906
    
    socket_server = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
    socket_server.bind((LOCAL_HOST, LOCAL_PORT))
    socket_server.listen()
    conn, address = socket_server.accept()
    print(f"收到了客户端的连接，客户端信息是 {address}")
    HOME = os.getcwd()

    path_x = [0, 200, 400, 600, 800]
    path_y = [0, 200, 500, 300, 400]
    # 规划速度：根据距离动态调整速度
    current_p_index = 0  # 当前目标点索引

    boat_data = dict()

    delta_h = 0.1
    
    # 非线性kt参数
    K=3.1429
    T=21.348
    deltam=6.8207e-2
    alpha=6.8595e-3
    
    # boat运动仿真类输入参数
    my_mmg = motionSim.BoatMotionSim(delta_h=delta_h, K=K, T=T, deltam=deltam, alpha=alpha)
    
    recorder = dict()
    recorder['exp_anle'] = list()
    recorder['path_x'] = list()
    recorder['path_y'] = list()
    recorder['yaw'] = list()
    recorder['u'] = list()
    recorder['rudder'] = list()
    recorder['y_error'] = list()

    try:
        x = np.array([-50.0, -50.0, 0.0, 0.0, 0.0, 0.0])# 初始无人艇位置
        u = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])# 初始无人艇速度
        current_p_index = 0
        old_rsp = 0
        while (1):
            if (current_p_index > len(path_x) - 1):
                break
            exp_p = (path_x[current_p_index], path_y[current_p_index])
            
            # 动态计算当前目标点的速度
            distance_to_target = np.sqrt((exp_p[0] - x[0]) ** 2 + (exp_p[1] - x[1]) ** 2)
            max_speed = 2.0  # 最大速度
            min_speed = 0.0  # 最小速度
            exp_speed = max(min_speed, min(max_speed, (distance_to_target / 30.0)))  # 根据距离动态调整速度

            if distance_to_target < 30:
                current_p_index += 1
                continue
            
            # 接收转速和舵角
            info, address = conn.recvfrom(1024)
            old_rsp = rsp
            delta, rsp = decode_data(info)

            # 计算更新 位置 速度
            x, u, nu, acc = my_mmg.runMMG(x=x, u=u, delta=delta * np.pi / 180.0, f=rsp)

            # position to exp_angle
            exp_angle = np.arctan2(exp_p[0] - x[0], exp_p[1] - x[1]) * 180.0 / np.pi

            # 为强化学习用
            observation = dict()
            observation['psi'] = rad_limit(rad_limit(x[2]) - exp_angle * np.pi / 180.0) * 180.0 / np.pi  # 角度偏差
            observation['dpsi'] = u[2] * 180.0 / np.pi  # Z轴角速度 弧度转角度
            observation['rud'] = delta  # 舵角 角度
            observation['dspeed'] = exp_speed - u[0]  # x速度偏差 船体坐标系下的轴向速度
            observation['v_x'] = u[0]  # 船体坐标系x速度
            observation['v_y'] = u[1]  # 船体坐标系y速度
            observation['acc_x'] = acc[0]
            observation['acc_n'] = acc[2] * 180.0 / np.pi
            observation['drsp'] = 0.01 * (rsp - old_rsp)  # 转速波动过大惩罚 收敛不需要转速归零
            observation['rsp'] = rsp * 0.01  # 转速

            boat_data['reward'] = get_reward(observation)
            boat_data['observation'] = observation
            boat_data['terminated'] = 0
            boat_data['truncated'] = 0
            data = encode_data(trans=boat_data)

            # 发送数据
            conn.send(data)

            # 计算横向偏差 点到直线距离
            if current_p_index > 0:
                y_error = distance_point_to_line(x[0], x[1], path_x[current_p_index], path_y[current_p_index], path_x[current_p_index - 1],
                                                 path_y[current_p_index - 1])
            else:
                y_error = np.sqrt((exp_p[0] - x[0]) ** 2 + (exp_p[1] - x[1]) ** 2)

            print(f'psi is {x[5]}, exp_angle is {exp_angle}')
            recorder['path_x'].append(x[0])
            recorder['path_y'].append(x[1])
            recorder['yaw'].append(x[5])
            recorder['u'].append(u)
            recorder['exp_anle'].append(exp_angle)
            recorder['rudder'].append(delta)
            recorder['y_error'].append(y_error)

            # 绘图pygame
            

        plt.figure()
        plt.plot(path_x, path_y)
        plt.plot(recorder['path_x'], recorder['path_y'])
        plt.figure()
        plt.plot(recorder['rudder'])
        plt.show()
        
    except Exception:
        # 关闭连接
        conn.close()
        socket_server.close()
        traceback.print_exc()

